#!/usr/bin/env python2
#coding:utf-8
"""
Author:  Rusty Stiles --<rstiles@gmail.com>
Purpose: Module to interact with remote server

TODO:

Write up on ssh/config

Write Doc strings, code error handling


Created: 06/21/2011
"""

import paraproxy, paramiko
import os
from stat import S_IRWXG,S_IRWXU


class Client:
    """SSH client to support operations on a remote machine"""

    def __init__(self, hostname, username, password=None):
        """Initialize SSH Client with the option to auto add host keys.
        Then connects to the remote machine, hostname, as username. You MUST be 
        able to connect to the machine in one 'hop', in order to connect through
        a proxy you must configure your .ssh/config""" 

        self.hostname = hostname
        self.username = username
        self.client = paramiko.SSHClient()
        self.client.load_system_host_keys()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(hostname,username=username,password=password)

    def get_pcpu_status(self, info_level=4):
        command = 'ps aux --sort -pcpu | head -n ' + str(info_level)
        info = self._get_stdout_of_command(command)
        header = info[0]
        info = info[1:]
        return header, info

    def _get_stdout_of_command(self, command):
        stdin, stdout, stderr = self.client.exec_command(command)
        info = [line.strip() for line in stdout.readlines()]
        error = [line.strip() for line in stderr.readlines()]
        if error:
            for line in error:
                print(line)
        return info

    def start_run(self, directory, code = './e.out', temp = 'temp', timed = False):
        program = os.path.join(directory, code)
        temp_file = os.path.join(directory, temp)
        if timed:
            command = "cd " + directory + " && ((time " + program + ") > " + temp_file + " 2>&1) &"
        else:
            command = "cd " + directory + " && " + program + " > " + temp_file + " &"
        self.client.exec_command(command)
    
    def remote_tail(self,filename, tail_level=4):
        command = 'tail -'+str(tail_level)+' '+filename
        info = self._get_stdout_of_command(command)
        return info
    
    def _ftp(func):
        def wrapper(self, *args, **kwargs):
            self.ftp = self.client.open_sftp()
            func(self, *args, **kwargs)
            self.ftp.close()
        return wrapper
    
    @_ftp
    def retrieve_file(self, remote_file, local_file, delete_remote=False):
        self.ftp.get(remote_file, local_file)
        if delete_remote:
            self.ftp.remove(remote_file)

    @_ftp
    def move(self, old_path, new_path, delete_remote=False):
        try:
            self.ftp.rename(old_path, new_path)
        except IOError:
            self.remote_mkdir(os.path.abspath(os.path.dirname(new_path)))
            self.ftp.rename(old_path, new_path)
        if delete_remote:
            self.ftp.remove(old_path)

    def remote_mkdir(self, folder):
        command = 'mkdir -p ' + folder
        print(folder + ' not on node ' + self.hostname + ' making it now.')
        self._get_stdout_of_command(command)

    def _ftp_put(self, ftp, folder, local_file, remote_file):
        try:
            ftp.put(local_file, remote_file)
        except IOError:
            self.remote_mkdir(folder)
            ftp.put(local_file, remote_file)

    @_ftp
    def put_file(self, local_file, remote_file, make_executable = False):
        folder = os.path.dirname(remote_file)
        self._ftp_put(self.ftp, folder, local_file, remote_file)
        if make_executable:
            self.ftp.chmod(remote_file, S_IRWXU | S_IRWXG)

    def compile_on_node(self, source_code, node_directory, compiled_name):
        """Compile the source_code on the specified node and returns the compiled code to the current same dir as
        the source code
        Returns """
        source_name = os.path.basename(source_code)
        remote_file = os.path.join(node_directory, source_name)
        returned_file = os.path.join(os.path.dirname(source_code), (source_name[:-2] + '.out'))
        compiled_path = os.path.join(node_directory, compiled_name)
        self.put_file(source_code, remote_file)
        command = 'ifort -o ' + os.path.join(node_directory, compiled_name) + ' ' + remote_file
        self._get_stdout_of_command(command)
        try:
            self.retrieve_file(compiled_path, returned_file)
            success = True
        except IOError:
            print("Error returning " + compiled_path + " Try again")
            success = False
        return success, returned_file

    def close(self):
        self.client.close()







def start_runs_on_nodes(nodes, work_dirs, username = 'stiles', code = 'e.out', temp_file = 'temp', timed = False):
    for node in nodes:
        client = Client(node, username)
        for folder in work_dirs:
            print("Starting run on " + node + " in " +folder)
            client.start_run(folder, code, temp_file, timed)
        client.close()


def compile_for_nodes(nodes, source_code, work_dirs, username = 'stiles', compiled_name = 'e.out'):
    """ **TEMP** use make_mode_list for nodes, get source path with os.path.abspath(input-source)
    work_directories is usually ['/work/<username>/', '/work/<username>/b/]
    We first compile the code on the first node and then copy it to all the dirs in work_dirs
    """
    client = Client(nodes[0], username)
    success, compiled_code = client.compile_on_node(source_code, work_dirs[0], compiled_name)
    client.close()

    if success:
        for node in nodes:
            client = Client(node, username)
            for direc in work_dirs:
                remote_path = os.path.join(direc, compiled_name)
                client.put_file(compiled_code, remote_path, make_executable = True)
            client.close()

def killall(nodes, username, process = 'e.out'):

    for node in nodes:
        client = Client(node,username)
        client.kill_process(process)
        client.close()

def get_progress(nodes, temp_files, username = 'stiles'):
    """The files that log the progress are made up of the work directory (/work), the user_dir on the nodes (stiles) and
    the actual file holding the progress are the first file and second file. the second file should be relative to the
    directory of the first file. so the defaults wold read the files in /work/stiles/temp and /work/stiles/b/temp
    """
    for host in nodes:
        client = Client(host,username)
        for fil in temp_files:
            info = client.remote_tail(fil)
            print('\n' + host + ', ' + fil + '\n')
            for line in info:
                print(line)
        client.close()


def make_node_list(cluster = 'e',first_node = 65, last_node = 88):
    """Cluster is the cluster that the nodes are on such as c,d,e,f... first node and last node are the number of the
    nodes that you want  e.g for nodes 'e65,e66,e67,e68' the cluster is 'e' and the first_node is 65 and the last node
    is 68."""
    last_node = int(last_node) + 1
    names = range(int(first_node), int(last_node))
    nodes = [cluster + str(name) for name in names]
    return nodes


def get_processor_status(nodes, username = 'stiles'):

    first = True

    for host in nodes:
        client = Client(host,username)
        header, info = client.get_pcpu_status()
        print('\n')
        if first:
            print(header + '\n')
            first = False
        print(host)
        for line in info:
            print(line)
        client.close()

def run_commands(nodes, command, username):
    for host in nodes:
        client = Client(host,username)
        print("Running " + command + " on " + host)
        info = client._get_stdout_of_command(command)
        for lin in info:
            print(lin)
        client.close()


if __name__=='__main__':
    pass
